#include "storage/disk_manager.h"
#include <cassert>
#include <iostream>
#include <string>
#include <cstring>
#include <vector>

// Test basic DiskManager functionality: create, read, write pages
void testBasicPageOperations() {
    std::cout << "Testing DiskManager basic page operations...\n";
    
    std::string filename = "test_disk_manager.db";
    std::remove(filename.c_str()); // Ensure no residual file before test starts
    
    // Create DiskManager
    DiskManager diskManager(filename);
    
    // Test allocatePage - allocate a new page
    uint32_t pageId = diskManager.allocatePage();
    assert(pageId == 0); // According to implementation, first page ID should be 0
    
    // Test writePage - write some data to the page
    char writeData[PAGE_SIZE] = {0};
    for (int i = 0; i < 100; ++i) {
        writeData[i] = static_cast<char>(i % 256);
    }
    diskManager.writePage(pageId, writeData);
    
    // Test readPage - read data from the page
    char readData[PAGE_SIZE] = {0};
    diskManager.readPage(pageId, readData);
    
    // Verify that the read data matches the written data
    for (int i = 0; i < 100; ++i) {
        assert(writeData[i] == readData[i]);
    }
    
    std::cout << "Basic page operations test passed!\n";
    std::remove(filename.c_str());
}

// Test DiskManager multiple page allocation and operations
void testMultiplePages() {
    std::cout << "Testing DiskManager multiple pages operations...\n";
    
    std::string filename = "test_disk_manager_multi.db";
    std::remove(filename.c_str());
    
    DiskManager diskManager(filename);
    
    // Allocate multiple pages
    const int numPages = 5;
    std::vector<uint32_t> pageIds;
    
    for (int i = 0; i < numPages; ++i) {
        uint32_t pageId = diskManager.allocatePage();
        pageIds.push_back(pageId);
        assert(pageId == i); // Verify page IDs are consecutive starting from 0
        
        // Write different data patterns to each page
        char writeData[PAGE_SIZE] = {0};
        for (int j = 0; j < 100; ++j) {
            writeData[j] = static_cast<char>((i * 100 + j) % 256);
        }
        diskManager.writePage(pageId, writeData);
    }
    
    // Verify all pages have correct data written
    for (int i = 0; i < numPages; ++i) {
        char readData[PAGE_SIZE] = {0};
        diskManager.readPage(pageIds[i], readData);
        
        for (int j = 0; j < 100; ++j) {
            assert(readData[j] == static_cast<char>((i * 100 + j) % 256));
        }
    }
    
    std::cout << "Multiple pages operations test passed!\n";
    std::remove(filename.c_str());
}

// Test DiskManager file size expansion functionality (ensureFileSize)
void testFileSizeExpansion() {
    std::cout << "Testing DiskManager file size expansion...\n";
    
    std::string filename = "test_disk_manager_expansion.db";
    std::remove(filename.c_str());
    
    DiskManager diskManager(filename);
    
    // Test allocating a non-consecutive page ID, which should trigger file size expansion
    uint32_t highPageId = 100; // Skip the first 99 pages
    
    char writeData[PAGE_SIZE] = {0};
    for (int i = 0; i < 100; ++i) {
        writeData[i] = static_cast<char>(i + 100);
    }
    
    // This should expand the file size to accommodate page ID 100
    diskManager.writePage(highPageId, writeData);
    
    // Verify that the written data can be correctly read
    char readData[PAGE_SIZE] = {0};
    diskManager.readPage(highPageId, readData);
    
    for (int i = 0; i < 100; ++i) {
        assert(readData[i] == static_cast<char>(i + 100));
    }
    
    std::cout << "File size expansion test passed!\n";
    std::remove(filename.c_str());
}

// Test DiskManager page deallocation functionality
void testPageDeallocation() {
    std::cout << "Testing DiskManager page deallocation...\n";
    
    std::string filename = "test_disk_manager_dealloc.db";
    std::remove(filename.c_str());
    
    DiskManager diskManager(filename);
    
    // Allocate three pages
    uint32_t pageId1 = diskManager.allocatePage();
    uint32_t pageId2 = diskManager.allocatePage();
    uint32_t pageId3 = diskManager.allocatePage();
    
    assert(pageId1 == 0);
    assert(pageId2 == 1);
    assert(pageId3 == 2);
    
    // Deallocate the second page
    diskManager.deallocatePage(pageId2);
    
    // Allocate a page again, it should reuse the deallocated page ID
    uint32_t newPageId = diskManager.allocatePage();
    assert(newPageId == pageId2); // Should reuse ID 1
    
    // Verify that the newly allocated page data is correctly initialized
    char readData[PAGE_SIZE] = {0};
    diskManager.readPage(newPageId, readData);
    
    // Verify page is initialized to all zeros
    bool allZeros = true;
    for (int j = 0; j < PAGE_SIZE; ++j) {
        if (readData[j] != 0) {
            allZeros = false;
            break;
        }
    }
    assert(allZeros);
    
    std::cout << "Page deallocation test passed!\n";
    std::remove(filename.c_str());
}

// Test DiskManager error handling
void testErrorHandling() {
    std::cout << "Testing DiskManager error handling...\n";
    
    std::string filename = "test_disk_manager_error.db";
    std::remove(filename.c_str());
    
    // Test scenario 1: Reading a non-existent page should throw an exception
    try {
        DiskManager diskManager(filename);
        char data[PAGE_SIZE] = {0};
        
        // Attempt to read an unallocated page, should throw exception
        diskManager.readPage(999, data);
        assert(false); // Test fails if this line is executed
    } catch (const std::runtime_error& e) {
        std::cout << "Caught expected exception: " << e.what() << std::endl;
    }
    
    // Test scenario 2: Behavior when file is already open
    { // Scope ensures diskManager1 is destroyed
        DiskManager diskManager1(filename);
        uint32_t pageId = diskManager1.allocatePage();
        
        char data[PAGE_SIZE] = {0};
        data[0] = 'A';
        diskManager1.writePage(pageId, data);
    }
    
    // Reopen the same file and verify data still exists
    { 
        DiskManager diskManager2(filename);
        char data[PAGE_SIZE] = {0};
        diskManager2.readPage(0, data);
        assert(data[0] == 'A');
    }
    
    std::cout << "Error handling test passed!\n";
    std::remove(filename.c_str());
}

// Test DiskManager file persistence
void testFilePersistence() {
    std::cout << "Testing DiskManager file persistence...\n";
    
    std::string filename = "test_disk_manager_persistence.db";
    std::remove(filename.c_str());
    
    // Create DiskManager, write data, then destroy it
    { 
        DiskManager diskManager(filename);
        uint32_t pageId = diskManager.allocatePage();
        
        char writeData[PAGE_SIZE] = {0};
        for (int i = 0; i < 100; ++i) {
            writeData[i] = static_cast<char>(i + 50);
        }
        diskManager.writePage(pageId, writeData);
        // diskManager should automatically close file when destroyed
    }
    
    // Recreate DiskManager, verify data still exists
    { 
        DiskManager diskManager(filename);
        char readData[PAGE_SIZE] = {0};
        diskManager.readPage(0, readData);
        
        for (int i = 0; i < 100; ++i) {
            assert(readData[i] == static_cast<char>(i + 50));
        }
    }
    
    std::cout << "File persistence test passed!\n";
    std::remove(filename.c_str());
}

int main() {
    std::cout << "Running DiskManager unit tests...\n\n";
    
    // Run all tests
    testBasicPageOperations();
    std::cout << "\n";
    
    testMultiplePages();
    std::cout << "\n";
    
    testFileSizeExpansion();
    std::cout << "\n";
    
    testPageDeallocation();
    std::cout << "\n";
    
    testErrorHandling();
    std::cout << "\n";
    
    testFilePersistence();
    
    std::cout << "\nAll DiskManager unit tests passed!\n";
    return 0;
}